import os
import shutil
from unittest import runner
try:
    from .cvio import cvio
    from .running_faster import Runner
except:
    from cvio import cvio
    from running_faster import Runner

def parse_text(text):
    if not isinstance(text, str):
        return text
    if os.path.isfile(text):
        assert text.endswith('.txt')
        with open(text) as fp:
            labels = [x.strip().replace('\t', '').replace('\n', '') for x in fp.readlines()]
        return labels
    labels = text.replace('\t', ',').replace('\n', ',').replace('，', ',').replace(' ', ',').split(',')
    return labels

def sample_one_img(i, n, src, dst, label):
    print('[%d/%d]' % (i, n), src)
    dst = os.path.join(dst, label)
    if not os.path.exists(dst):
        os.makedirs(dst)
    dst = os.path.join(dst, os.path.basename(src))
    if os.path.exists(dst):
        return
    shutil.copy(src, dst)

def sample_scls_data_runner(src, dst, label, worker='process', num_workers=8, spec=''):
    labels = parse_text(label)
    # print(labels)
    runner = Runner()
    images = cvio.load_image_list(src, recursive=True, silent=False)
    new_images = []
    for img in images:
        if spec not in img:
            continue

        for label in labels:
            if label in img:
                if os.path.exists(os.path.join(dst, label, os.path.basename(img))):
                    continue
                new_images.append(img)
                break

    n = len(new_images)
    for i, img in enumerate(new_images, 1):
        label = img.split(os.sep)[-2]
        runner.append(sample_one_img, (i, n, img, dst, label))

    runner.run(mode=worker, num_workers=num_workers)

if __name__ == '__main__':
    src = r'G:\data\scls_data\drink'
    dst = r'G:\data\scls_data\sample_test_0714'
    label = ''
    sample_scls_data_runner(src, dst, label, worker='process', num_workers=8)